// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

// Benchmark http-server by multiple threads.

#include <krpc/kthread/kthread.h>
#include <krpc/utility/logging.h>
#include <krpc/rpc/channel.h>
#include <krpc/rpc/server.h>
#include <krpc/var/kvar.h>
#include <krpc/rpc/restful_client.h>

TURBO_FLAG(std::string, data, "", "POST this data to the http server");
TURBO_FLAG(int32_t, thread_num, 50, "Number of threads to send requests");
TURBO_FLAG(bool, use_kthread, false, "Use kthread to send requests");
TURBO_FLAG(std::string, connection_type, "", "Connection type. Available values: single, pooled, short");
TURBO_FLAG(std::string, url, "http://0.0.0.0:8010/HttpService/Echo", "url of server");
TURBO_FLAG(std::string, load_balancer, "", "The algorithm for load balancing");
TURBO_FLAG(int32_t, timeout_ms, 100, "RPC timeout in milliseconds");
TURBO_FLAG(int32_t, max_retry, 3, "Max retries(not including the first RPC)"); 
TURBO_FLAG(bool, dont_fail, false, "Print fatal when some call failed");
TURBO_FLAG(int32_t, dummy_port, -1, "Launch dummy server at this port");
TURBO_FLAG(std::string, protocol, "http", "Client-side protocol");

kvar::LatencyRecorder g_latency_recorder("client");

static void* sender(void* arg) {
    krpc::Channel* channel = static_cast<krpc::Channel*>(arg);

    while (!krpc::IsAskedToQuit()) {
        // We will receive response synchronously, safe to put variables
        // on stack.
        krpc::Controller cntl;

        cntl.set_timeout_ms(turbo::get_flag(FLAGS_timeout_ms)/*milliseconds*/);
        cntl.set_max_retry(turbo::get_flag(FLAGS_max_retry));
        cntl.http_request().uri() = turbo::get_flag(FLAGS_url);
        if (!turbo::get_flag(FLAGS_data).empty()) {
            cntl.http_request().set_method(krpc::HTTP_METHOD_POST);
            cntl.request_attachment().append(turbo::get_flag(FLAGS_data));
        }

        // Because `done'(last parameter) is NULL, this function waits until
        // the response comes back or error occurs(including timedout).
        channel->CallMethod(NULL, &cntl, NULL, NULL, NULL);
        if (!cntl.Failed()) {
            g_latency_recorder << cntl.latency_us();
        } else {
            CHECK(krpc::IsAskedToQuit() || !turbo::get_flag(FLAGS_dont_fail))
                << "error=" << cntl.ErrorText() << " latency=" << cntl.latency_us();
            // We can't connect to the server, sleep a while. Notice that this
            // is a specific sleeping to prevent this thread from spinning too
            // fast. You should continue the business logic in a production 
            // server rather than sleeping.
            kthread_usleep(100000);
        }
    }
    return NULL;
}

int main(int argc, char* argv[]) {

    auto *run_app = turbo::Servlet::instance().run_app();
    run_app->add_option("--data", FLAGS_data, FLAGS_data.help());
    run_app->add_option("--thread", FLAGS_thread_num, FLAGS_thread_num.help());
    run_app->add_option("--url", FLAGS_url, FLAGS_url.help());
    run_app->add_option("--protocol", FLAGS_protocol, FLAGS_protocol.help());
    run_app->add_option("--dummy_port", FLAGS_dummy_port, FLAGS_dummy_port.help());
    run_app->add_option("--max_retry", FLAGS_max_retry, FLAGS_max_retry.help());
    run_app->add_option("--timeout", FLAGS_timeout_ms, FLAGS_timeout_ms.help());
    run_app->add_option("--lb", FLAGS_load_balancer, FLAGS_load_balancer.help());
    run_app->add_option("--use_kthread", FLAGS_use_kthread, FLAGS_use_kthread.help());
    run_app->add_option("--connection_type", FLAGS_connection_type, FLAGS_connection_type.help());
    run_app->add_option("--dont_fail", FLAGS_dont_fail, FLAGS_dont_fail.help());
    auto [exit, ret] = turbo::Servlet::instance().run(argc, argv);
    if (exit) {
        return ret;
    }

    // A Channel represents a communication line to a Server. Notice that 
    // Channel is thread-safe and can be shared by all threads in your program.
    krpc::Channel channel;
    krpc::ChannelOptions options;
    options.protocol = turbo::get_flag(FLAGS_protocol);
    options.connection_type = turbo::get_flag(FLAGS_connection_type);
    auto r = krpc::parse_endpoint(turbo::get_flag(FLAGS_url));
    if(!r.ok()) {
        LOG(ERROR) << "Fail to parse host from url";
        return -1;
    }
    // turbo::get_flag(FLAGS_load_balancer).c_str(),
    if (channel.Init(r.value(),  &options) != 0) {
        LOG(ERROR) << "Fail to initialize channel";
        return -1;
    }

    std::vector<kthread_t> bids;
    std::vector<pthread_t> pids;
    if (!turbo::get_flag(FLAGS_use_kthread)) {
        pids.resize(turbo::get_flag(FLAGS_thread_num));
        for (int i = 0; i < turbo::get_flag(FLAGS_thread_num); ++i) {
            if (pthread_create(&pids[i], NULL, sender, &channel) != 0) {
                LOG(ERROR) << "Fail to create pthread";
                return -1;
            }
        }
    } else {
        bids.resize(turbo::get_flag(FLAGS_thread_num));
        for (int i = 0; i < turbo::get_flag(FLAGS_thread_num); ++i) {
            if (kthread_start_background(
                    &bids[i], NULL, sender, &channel) != 0) {
                LOG(ERROR) << "Fail to create kthread";
                return -1;
            }
        }
    }

    if (turbo::get_flag(FLAGS_dummy_port) >= 0) {
        krpc::StartDummyServerAt(turbo::get_flag(FLAGS_dummy_port));
    }

    while (!krpc::IsAskedToQuit()) {
        sleep(1);
        LOG(INFO) << "Sending " << turbo::get_flag(FLAGS_protocol) << " requests at qps="
                  << g_latency_recorder.qps(1)
                  << " latency=" << g_latency_recorder.latency(1);
    }

    LOG(INFO) << "benchmark_http is going to quit";
    for (int i = 0; i <turbo::get_flag( FLAGS_thread_num); ++i) {
        if (!turbo::get_flag(FLAGS_use_kthread)) {
            pthread_join(pids[i], NULL);
        } else {
            kthread_join(bids[i], NULL);
        }
    }

    return 0;
}
